#include "concurrent.hpp"

namespace libco {

extern TaskQueue* s_global_q;

constexpr uint8_t MAX_LOCK_TRIES = 4;
constexpr uint8_t MAX_NORMAL_WAITING_MS = 10;


TaskQueue::TaskQueue(size_t limit) {
    m_limit = (limit == 0) ? (size_t(-1)) : limit;
}

void TaskQueue::push(Coroutine::ptr elem) {
    Lock lock(m_mutex);
    m_cv_notfull.wait(lock, [this](){
        return m_que.size() < m_limit;
    });
    m_que.push_back(elem);
    m_cv_notempty.notify_one();
}

void TaskQueue::push_front(Coroutine::ptr elem) {
    Lock lock(m_mutex);
    m_cv_notfull.wait(lock, [this](){
        return m_que.size() < m_limit;
    });
    m_que.push_front(elem);
    m_cv_notempty.notify_one();
}

Coroutine::ptr TaskQueue::pop() {
    Lock lock(m_mutex);
    m_cv_notempty.wait(lock, [this](){
        return !m_que.empty();
    });
    Coroutine::ptr top_elem = std::move(m_que.front());
    m_que.pop_front();
    m_cv_notfull.notify_one();
    return top_elem;
}

std::optional<Coroutine::ptr> TaskQueue::top() {
    Lock lock(m_mutex);
    if (m_que.empty()) {
        return std::nullopt;
    }
    Coroutine::ptr top_elem = m_que.front();
    return top_elem;
}


bool TaskQueue::try_push(Coroutine::ptr elem) {
    Lock lock(m_mutex);
    if (m_que.size() >= m_limit) {
        return false;
    }
    m_que.push_back(std::move(elem));
    m_cv_notempty.notify_one();
    return true;
}

size_t TaskQueue::try_push_n(std::vector<Coroutine::ptr>& elems, size_t n) {
    Lock lock(m_mutex);
    size_t success = 0;
    while (m_que.size() < m_limit && success < n) {
        m_que.push_back(std::move(elems[success]));
        ++success;
    }
    m_cv_notempty.notify_one();
    return success;
}

void TaskQueue::push_all(std::vector<Coroutine::ptr>& elems) {
    Lock lock(m_mutex);
    size_t n = elems.size(), success = 0;
    while(success < n) {
        while (m_que.size() < m_limit) {
            m_que.push_back(std::move(elems[success]));
            ++success;
        }
    }
    m_cv_notempty.notify_one();
}

size_t TaskQueue::try_pop_n(std::vector<Coroutine::ptr>& elems, size_t n) {
    Lock lock(m_mutex);
    size_t success = 0;
    while (!m_que.empty() && success < n) {
        elems.push_back(m_que.front());
        m_que.pop_front();
        ++success;
    }
    m_cv_notfull.notify_one();
    return success;
}

size_t TaskQueue::try_pop_all(std::vector<Coroutine::ptr>& elems) {
    Lock lock(m_mutex);
    size_t success = 0;
    while (!m_que.empty()) {
        elems.push_back(m_que.front());
        m_que.pop_front();
        ++success;
    }
    m_cv_notfull.notify_one();
    return success;
}

size_t TaskQueue::try_push_all(std::vector<Coroutine::ptr>& elems) {
    Lock lock(m_mutex);
    size_t success = 0;
    while (m_que.size() < m_limit && success < elems.size()) {
        m_que.push_back(elems[success]);
        ++success;
    }
    m_cv_notempty.notify_one();
    return success;
}

std::optional<Coroutine::ptr> TaskQueue::try_pop() {
    Lock lock(m_mutex);
    if (m_que.empty()) {
        return std::nullopt;
    }
    Coroutine::ptr top_elem = m_que.front();
    m_que.pop_front();
    m_cv_notfull.notify_one();
    return top_elem;
}

void TaskQueue::lock() { m_mutex.lock(); }
void TaskQueue::unlock() { m_mutex.unlock(); }
size_t TaskQueue::tasks() { Mutex::Lock lock(m_mutex) ; return m_que.size(); }


} // namespace libco

namespace libco {
    
CoMutex::CoMutex() {
    m_locked.clear();
}

void CoMutex::lock() {
    while (!m_locked.test_and_set()) {
        auto cur = libco::Coroutine::GetThis();
        m_waiting_que.push(cur);
        cur->yield();
    }
}

// void CoMutex::lock() {
//     // 加锁成功，返回
//     if (!m_locked.test_and_set()) {
//         return;
//     }
//     if (m_starving) {
//         auto cur = libco::Coroutine::GetThis();
//         m_waiting_que.push(cur);
//         cur->yield();
//     }
//     lockslow();
// }

// void CoMutex::lockslow() {
//     uint8_t tries = 0;
//     uint64_t start_time = GetCurrentTimeMs();
    
//     while(m_locked.test_and_set()) {
//         if (!m_starving) {
//             ++tries;
//             if (tries >= MAX_LOCK_TRIES) {
//                 auto cur = libco::Coroutine::GetThis();
//                 m_waiting_que.push(cur);
//                 cur->yield();
//                 LIBCO_LOG_DEBUG << "coroutine back to get lock ";
//                 if (!m_starving && GetCurrentTimeMs() - start_time >= MAX_NORMAL_WAITING_MS) {
//                     m_starving = true;
//                     continue;
//                 } else {
//                     m_waiting_que.push(cur);
//                     cur->yield();
//                     tries = 0;
//                 }
//             }
//         } else {
//             LIBCO_LOG_DEBUG << "coroutine is starving";
//             auto cur = libco::Coroutine::GetThis();
//             m_waiting_que.push_front(cur);
//             cur->yield();
//         }
//     }
//     // 恢复正常模式
//     if(m_starving) {
//         LIBCO_LOG_DEBUG << "coroutine get lock while starving";
//         m_starving = false;
//     }
// }

void CoMutex::unlock() {
    m_locked.clear();
    auto front = m_waiting_que.try_pop();
    if (front.has_value()) {
        s_global_q->push(front.value());
    }
}
    
} // namespace libco

namespace libco {

void WaitGroup::done() {   
    Mutex::Lock lock(m_mutex);
    --m_count;
    LIBCO_ASSERT(m_count >= 0);
    if (m_count == 0) {
        m_cv.notify_all();
    }
}

void WaitGroup::add() {
    Mutex::Lock lock(m_mutex);
    ++m_count;  
}

void WaitGroup::wait() {
    Mutex::Lock lock(m_mutex);
    m_cv.wait(lock, [this](){
        return m_count == 0;
    });
}

} // namespace libco